# -*- coding:utf-8 -*-
# author:jmz
 # gevent 实现单线程下的socket并发，一样要放在导入socket模块之前。。。。。。天呢，浪费了一天的时间才找到原因
from gevent import spawn, joinall, monkey
monkey.patch_all()

from socket import *
import hashlib
import base64
import struct



class Worker(object):
    def __init__(self, conn):
        self.conn = conn

    def open(self, data):
        self.send(data + 'sb')

    def message(self, data):
        self.send(data + 'sb')

    def send(self, data):
        send_msg(self.conn, data)

    def close(self):
        self.conn.close()


def send_msg(conn, data):
    """
     WebSocket服务端向客户端发送消息
    :param conn: 客户端连接到服务器端的socket对象,即： conn,address = socket.accept()
    :param data: 向客户端发送的str 数据
    """
    msg_bytes = bytes(data, encoding="utf-8")
    token = b"\x81"
    length = len(msg_bytes)
    if length < 126:
        token += struct.pack("B", length)
    elif length <= 0xFFFF:
        token += struct.pack("!BH", 126, length)
    else:
        token += struct.pack("!BQ", 127, length)
    msg = token + msg_bytes
    conn.send(msg)


def get_msg(conn):
    """
    这是websocet 用于接收处理客户端发过来的数据
    :param conn:
    :return: str
    """
    info = conn.recv(8096)  # 接收的是帧数据
    payload_len = info[1] & 127
    if payload_len == 126:
        extend_payload_len = info[2:4]
        mask = info[4:8]
        decoded = info[8:]
    elif payload_len == 127:
        extend_payload_len = info[2:10]
        mask = info[10:14]
        decoded = info[14:]
    else:
        extend_payload_len = None
        mask = info[2:6]
        decoded = info[6:]
    bytes_list = bytearray()

    for i in range(len(decoded)):
        chunk = decoded[i] ^ mask[i % 4]
        bytes_list.append(chunk)
    msg = str(bytes_list, encoding='utf-8')
    return msg


def run(conn):
    worker = Worker(conn)
    try:
        msg = get_msg(worker.conn)
        worker.open(msg)
        while True:
            msg = get_msg(worker.conn)
            if not msg:
                worker.close()
                return False
            worker.message(msg)
    except:
        worker.close()

class Websocket(object):
    magic = '258EAFA5-E914-47DA-95CA-C5AB0DC85B11'
    coding = "utf-8"
    response = "HTTP/1.1 101 Switching Protocols\r\n" \
               "Upgrade:websocket\r\n" \
               "Connection: Upgrade\r\n" \
               "Sec-WebSocket-Accept: {1}\r\n" \
               "WebSocket-Location: ws://{2}/chat\r\n" \
               "WebSocket-Protocol: chat\r\n\r\n"

    def __init__(self, host, addr, listen_num=5):
        '''
        确认 ip 和 端口
        :param host:str
        :param addr:int
        '''
        self.host = host
        self.addr = addr
        self.listen_num = listen_num
        self.ws = socket(AF_INET, SOCK_STREAM)
        self.ws.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
        self.ws.bind((self.host, self.addr))
        self.ws.listen(self.listen_num)

    def run(self):
        while True:
            if (self._connect()):
                spawn(run, self.conn)

    def _connect(self):
        self.conn, _ = self.ws.accept()
        try:
            data = self.conn.recv(8129)
            if not data:
                self.conn.close()
                return False
            data = data.decode(encoding="utf-8")
            for line in data.split("\r\n\r\n")[0].split("\r\n")[1:]:
                k, v = line.split(": ")
                if k == 'Sec-WebSocket-Key':
                    sha1 = hashlib.sha1()
                    sha1.update((v + self.magic).encode(encoding=self.coding))
                    key = base64.b64encode(sha1.digest()).decode(encoding=self.coding)
                    self.conn.send(
                        self.response.replace('{1}', key).replace('{2}', "%s:%s" % (self.host, self.addr)).encode(
                            encoding=self.coding))
                    return True
            self.conn.close()
            return False
        except Exception as e:
            print(e)
            self.conn.close()
            return False


if __name__ == '__main__':
    ws = Websocket('127.0.0.1', 3389, 5)
    ws.run()
